'''
对三分类数据集（腺癌，鳞癌，其他）进行分布统计
绘制统计散点图

'''


import matplotlib.pyplot as plt
import pandas as pd
import os

# 读取统计数据
def read_data():
    csv_path = 'D:/lung_cancer/data/divide_csv/three/train.csv'
    data = pd.read_csv(csv_path)
    return data

# 将某类数据按类别划分
def devide_data(data, cancer_type):

    x_list = [[],[]]
    x_list[0].extend(data)
    x_list[1].extend(cancer_type)

    return x_list


def draw_box(one_data, name):

    son = name.split('/')[-1]
    print('------start draw ', son, '------')

    fig, axs = plt.subplots(2, figsize=(10,10))
    axs[0].scatter(one_data[0], one_data[1], s=2)  # x横坐标，y纵坐标
    axs[0].set(title='scatter_plot '+son)
    axs[1].hist2d(one_data[0], one_data[1])
    axs[1].set(title='hist2d_plot '+son)

    # fig.savefig(name)
    plt.show()
    plt.close()

# 绘制三分类坐标分布
def draw_coord(type_list, x_list, y_list, name):

    print('------start draw coord------')
    for i in range(len(type_list)):
        if int(type_list[i]) == 1:
            plt.plot([x_list[i]], [y_list[i]], marker='o', color='green', markersize=2)
        elif int(type_list[i]) == 2:
            plt.plot([x_list[i]], [y_list[i]], marker='o', color='red', markersize=2)
        else:
            plt.plot([x_list[i]], [y_list[i]], marker='o', color='blue', markersize=2)

    plt.savefig(name)
    plt.show()
    plt.close()

# 对不同特征做统计
def statistic():
    data = read_data()
    cancer_type = data['cancer_type']

    save_path = 'D:/lung_cancer/lung_cancer_code/release_code/output/three/Scatterplot/'
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    # 统计病灶 X
    cancer_x = data['x']
    x_list = devide_data(cancer_x, cancer_type)
    draw_box(x_list, save_path + '/x')

    # 统计病灶 Y
    cancer_y = data['y']
    y_list = devide_data(cancer_y, cancer_type)
    draw_box(y_list, save_path + '/y')

    # 统计病灶 newX
    cancer_newX = data['cx']
    newx_list = devide_data(cancer_newX, cancer_type)
    draw_box(newx_list, save_path + '/cx')

    # 统计病灶 newY
    cancer_newY = data['cy']
    newy_list = devide_data(cancer_newY, cancer_type)
    draw_box(newy_list, save_path + '/cy')

    # 统计病灶 Z
    cancer_z = data['z']
    y_list = devide_data(cancer_z, cancer_type)
    draw_box(y_list, save_path + '/z')


    # 统计病灶半径
    cancer_r = data['r']
    r_list = devide_data(cancer_r, cancer_type)
    draw_box(r_list, save_path+'/r')

    # suv_max
    suvmax = data['global_suvmax']
    suvmax_list = devide_data(suvmax, cancer_type)
    draw_box(suvmax_list, save_path+'/suv_max')

    # suv_min
    suvmin = data['local_suvmin']
    suvmin_list = devide_data(suvmin, cancer_type)
    draw_box(suvmin_list, save_path+'/suv_min')

    # suv_avg
    suvavg = data['local_suvavg']
    suvavg_list = devide_data(suvavg, cancer_type)
    draw_box(suvavg_list, save_path+'/suv_avg')

    # suv_std
    suvstd = data['local_suvstd']
    suvstd_list = devide_data(suvstd, cancer_type)
    draw_box(suvstd_list, save_path+'/suv_std')

    # suv_var
    suvstd = data['local_suvvar']
    suvstd_list = devide_data(suvstd, cancer_type)
    draw_box(suvstd_list, save_path+'/suv_var')

    # age
    age = data['patientAge']
    age_list = devide_data(age, cancer_type)
    draw_box(age_list, save_path+'/age')

    # size
    age = data['patientSize']
    age_list = devide_data(age, cancer_type)
    draw_box(age_list, save_path + '/size')

    # weight
    patientWeight = data['patientWeight']
    patientWeight_list = devide_data(patientWeight, cancer_type)
    draw_box(patientWeight_list, save_path+'/patientWeight')

    # Sex
    patientSex = data['patientSex']
    patientSex_list = devide_data(patientSex, cancer_type)
    draw_box(patientSex_list, save_path+'/patientSex')

    # 统计坐标
    draw_coord(cancer_type, cancer_newX, cancer_newY, save_path + '/coord')

if __name__ == '__main__':

    statistic()